#!/usr/bin/env python
# coding: utf-8

# In[2]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os


TOKEN = "1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c"
ts.set_token(TOKEN)
pro = ts.pro_api()

# 设置保存图片的路径
save_path = "C:/Users/Administrator/Desktop/"

# 确保保存路径存在
if not os.path.exists(save_path):
    os.makedirs(save_path)

# 获取股票数据
def get_stock_data(symbol, start_date, end_date):
    df = pro.daily(ts_code=symbol, start_date=start_date, end_date=end_date)
    df.index = pd.to_datetime(df['trade_date'])
    df.sort_index(inplace=True)
    return df

# 交易信号生成
def generate_signals(df, short_window=20, long_window=50):
    df['MA_short'] = df['close'].rolling(window=short_window).mean()
    df['MA_long'] = df['close'].rolling(window=long_window).mean()
    df['signal'] = np.where(df['MA_short'] > df['MA_long'], 1, -1)
    return df

# 回测
def backtest(df, initial_capital=100000):
    df['position'] = df['signal'].shift(1)
    df['strategy_return'] = df['position'] * df['close'].pct_change()
    df['cumulative_return'] = (1 + df['strategy_return']).cumprod()
    df['portfolio_value'] = df['cumulative_return'] * initial_capital
    return df

# 绘制 K 线图和交易信号
def plot_kline_and_signals(df, short_window, long_window, symbol):
    plt.figure(figsize=(14, 7))
    plt.plot(df.index, df['close'], label='Close Price')
    plt.plot(df.index, df['MA_short'], label=f'MA{short_window}')
    plt.plot(df.index, df['MA_long'], label=f'MA{long_window}')
    plt.scatter(df[df['signal'] == 1].index, df[df['signal'] == 1]['close'], color='green', label='Buy Signal')
    plt.scatter(df[df['signal'] == -1].index, df[df['signal'] == -1]['close'], color='red', label='Sell Signal')
    plt.legend()
    plt.title(f"Kline and Trading Signals for {symbol}")
    plt.savefig(os.path.join(save_path, f"kline_signals_{symbol}.png"))  # 保存图片
    plt.close()

# 绘制收益率曲线
def plot_return_curve(df, symbol):
    plt.figure(figsize=(14, 7))
    plt.plot(df.index, df['cumulative_return'], label='Strategy Return')
    plt.legend()
    plt.title(f"Cumulative Return Curve for {symbol}")
    plt.savefig(os.path.join(save_path, f"return_curve_{symbol}.png"))  # 保存图片
    plt.close()

# 多支股票回测
def multi_stock_backtest(symbols, start_date, end_date):
    results = {}
    for symbol in symbols:
        df = get_stock_data(symbol, start_date, end_date)
        df = generate_signals(df, short_window=20, long_window=50)
        df = backtest(df)
        results[symbol] = df
    return results

# 比较收益率曲线
def compare_return_curves(results):
    plt.figure(figsize=(14, 7))
    for symbol, df in results.items():
        plt.plot(df.index, df['cumulative_return'], label=symbol)
    plt.legend()
    plt.title("Comparison of Cumulative Return Curves")
    plt.savefig(os.path.join(save_path, "comparison_return_curves.png"))  # 保存图片
    plt.close()

# 示例：获取平安银行（000001.SZ）从2024-01-01到2024-12-31的日线数据
df = get_stock_data(symbol="000001.SZ", start_date="20240101", end_date="20241231")

# 应用交易策略并可视化
df = generate_signals(df, short_window=20, long_window=50)
df = backtest(df)
plot_kline_and_signals(df, short_window=20, long_window=50, symbol="000001.SZ")
plot_return_curve(df, symbol="000001.SZ")

# 多支股票回测
symbols = ["000001.SZ", "600036.SH"]
results = multi_stock_backtest(symbols, start_date="20240101", end_date="20241231")
compare_return_curves(results)


# In[ ]:




